{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "99eee7a2",
   "metadata": {},
   "source": [
    "# MCP ToolSpec\n",
    "\n",
    "This tool connects to MCP Servers and allows an Agent to call the tools provided by MCP Servers.\n",
    "\n",
    "This idea is migrated from [Integrate MCP tools into LlamaIndex](https://psiace.me/posts/integrate-mcp-tools-into-llamaindex/).\n",
    "\n",
    "To run this example, you need to edit the .env file to set the correct values for the [IPInfo API Token](https://ipinfo.io/signup) and a [OpenAI API key](https://platform.openai.com/api-keys).\n",
    "\n",
    "```bash\n",
    "# get the code\n",
    "git clone https://github.com/run-llama/llama_index\n",
    "cd llama_index/llama-index-integrations/tools/llama-index-tools-mcp/examples\n",
    "\n",
    "# install dependencies\n",
    "pip install ipinfo\n",
    "\n",
    "cp .env.example .env\n",
    "# NOTE: edit the .env file to have the correct values!\n",
    "\n",
    "# run the server\n",
    "python mcp_server.py --server_type=sse\n",
    "```\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d8f6ef6b",
   "metadata": {},
   "source": [
    "In this example, we will create a toy example with an agent that can query the ip info of the user.\n",
    "\n",
    "It's built using the `AgentWorkflow` class from LlamaIndex. If that's new to you, you can [read more about it](https://docs.llamaindex.ai/en/stable/examples/agent/agent_workflow_basic/)!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30e6088f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.llms.openai import OpenAI\n",
    "import dotenv\n",
    "\n",
    "dotenv.load_dotenv()\n",
    "\n",
    "llm = OpenAI(model=\"gpt-4o\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f519717e-c620-4e32-9729-99740b648556",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.tools.mcp import McpToolSpec\n",
    "from llama_index.core.agent.workflow import FunctionAgent, ToolCallResult, ToolCall\n",
    "from llama_index.core.workflow import Context\n",
    "\n",
    "SYSTEM_PROMPT = \"\"\"\\\n",
    "You are an AI assistant.\n",
    "\n",
    "Before you help a user, you need to fetch the ip info first, to help you follow the laws of the country.\n",
    "\"\"\"\n",
    "\n",
    "\n",
    "async def get_agent(tools: McpToolSpec):\n",
    "    tools = await tools.to_tool_list_async()\n",
    "    agent = FunctionAgent(\n",
    "        name=\"Agent\",\n",
    "        description=\"An agent that can fetch the ip info of the user.\",\n",
    "        tools=tools,\n",
    "        llm=llm,\n",
    "        system_prompt=SYSTEM_PROMPT,\n",
    "    )\n",
    "    return agent\n",
    "\n",
    "\n",
    "async def handle_user_message(\n",
    "    message_content: str,\n",
    "    agent: FunctionAgent,\n",
    "    agent_context: Context,\n",
    "    verbose: bool = False,\n",
    "):\n",
    "    handler = agent.run(message_content, ctx=agent_context)\n",
    "    async for event in handler.stream_events():\n",
    "        if verbose and type(event) == ToolCall:\n",
    "            print(f\"Calling tool {event.tool_name} with kwargs {event.tool_kwargs}\")\n",
    "        elif verbose and type(event) == ToolCallResult:\n",
    "            print(f\"Tool {event.tool_name} returned {event.tool_output}\")\n",
    "\n",
    "    response = await handler\n",
    "    return str(response)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9a69bdec",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.tools.mcp import BasicMCPClient, McpToolSpec\n",
    "\n",
    "# We consider there is a mcp server running on 127.0.0.1:8000, or you can use the mcp client to connect to your own mcp server.\n",
    "mcp_client = BasicMCPClient(\"http://127.0.0.1:8000/sse\")\n",
    "mcp_tool = McpToolSpec(client=mcp_client)\n",
    "\n",
    "# get the agent\n",
    "agent = await get_agent(mcp_tool)\n",
    "\n",
    "# create the agent context\n",
    "agent_context = Context(agent)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "db33a67f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "User:  Hello! I need some help with my homework\n",
      "Calling tool fetch_ipinfo with kwargs {'kwargs': ''}\n",
      "Tool fetch_ipinfo returned meta=None content=[TextContent(type='text', text='{\"ip\": \"71.17.163.153\", \"hostname\": \"71-17-163-153.regn.hsdb.sasknet.sk.ca\", \"city\": \"Regina\", \"region\": \"Saskatchewan\", \"country\": \"CA\", \"loc\": \"50.4501,-104.6178\", \"timezone\": \"America/Regina\"}')] isError=False\n",
      "Agent:  Hello! How can I assist you with your homework?\n"
     ]
    }
   ],
   "source": [
    "# Run the agent!\n",
    "while True:\n",
    "    user_input = input(\"Enter your message: \")\n",
    "    if user_input == \"exit\":\n",
    "        break\n",
    "    print(\"User: \", user_input)\n",
    "    response = await handle_user_message(user_input, agent, agent_context, verbose=True)\n",
    "    print(\"Agent: \", response)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "79c44895",
   "metadata": {},
   "source": [
    "Here, we can see the agent is calling the `fetch_ipinfo` tool to get the ip info! This tool is running remotely on the mcp server.\n",
    "\n",
    "The `MCPToolSpec` is connecting to the MCP server and creating `FunctionTool`s for each tool that is registered on the MCP server."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7a84adf7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fetch_ipinfo Get the detailed information of a specified IP address\n",
      "\n",
      "    Args:\n",
      "        ip(str or None): The IP address to get information for. Follow the format like 192.168.1.1 .\n",
      "        **kwargs: Additional keyword arguments to pass to the IPInfo handler.\n",
      "    Returns:\n",
      "        IPDetails: The detailed information of the specified IP address.\n",
      "    \n"
     ]
    }
   ],
   "source": [
    "tools = await mcp_tool.to_tool_list_async()\n",
    "for tool in tools:\n",
    "    print(tool.metadata.name, tool.metadata.description)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "33e7be78",
   "metadata": {},
   "source": [
    "You can also limit the tools that the `MCPToolSpec` will create by passing a list of tool names to the `MCPToolSpec` constructor."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4929af70",
   "metadata": {},
   "outputs": [],
   "source": [
    "mcp_tool = McpToolSpec(client=mcp_client, allowed_tools=[\"some fake tool\"])\n",
    "tools = await mcp_tool.to_tool_list_async()\n",
    "for tool in tools:\n",
    "    print(tool.metadata.name, tool.metadata.description)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e9c9df25",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bb1bda99",
   "metadata": {},
   "source": [
    "**Alternatively**, \n",
    "\n",
    "You can directly use the `get_tools_from_mcp_url` or `aget_tools_from_mcp_url` functions to get a list of `FunctionTool`s from an MCP server."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e58cdcdd",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.tools.mcp import (\n",
    "    get_tools_from_mcp_url,\n",
    "    aget_tools_from_mcp_url,\n",
    ")\n",
    "\n",
    "# async\n",
    "tools = await aget_tools_from_mcp_url(\"http://127.0.0.1:8000/sse\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "74d6d50e",
   "metadata": {},
   "source": [
    "By default, this will use our `BasicMCPClient`, which will run a command or connect to the URL and return the tools.\n",
    "\n",
    "You can also pass in a custom `ClientSession` to use a different client.\n",
    "\n",
    "You can also specify a list of allowed tools to filter the tools that are returned."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7cda4721",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.tools.mcp import BasicMCPClient\n",
    "\n",
    "client = BasicMCPClient(\"http://127.0.0.1:8000/sse\")\n",
    "\n",
    "tools = await aget_tools_from_mcp_url(\n",
    "    \"http://127.0.0.1:8000/sse\",\n",
    "    client=client,\n",
    "    allowed_tools=[\"fetch_ipinfo\"],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "248fed05",
   "metadata": {},
   "source": [
    "Then create the agent directly using the obtained list of `FunctionTool`s."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a51077e1",
   "metadata": {},
   "outputs": [],
   "source": [
    "agent = FunctionAgent(\n",
    "    name=\"Agent\",\n",
    "    description=\"An agent that can fetch the ip info of the user.\",\n",
    "    tools=tools,\n",
    "    llm=llm,\n",
    "    system_prompt=SYSTEM_PROMPT,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d9524804",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run the agent!\n",
    "while True:\n",
    "    user_input = input(\"Enter your message: \")\n",
    "    if user_input == \"exit\":\n",
    "        break\n",
    "    print(\"User: \", user_input)\n",
    "    response = await handle_user_message(user_input, agent, agent_context, verbose=True)\n",
    "    print(\"Agent: \", response)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
